import gradio as gr
from PIL import Image
import numpy as np
from sklearn.cluster import KMeans

def extract_color_palette(image_data, num_colors=5):
    # 将numpy数组转换为PIL图像
    image = Image.fromarray(image_data)
    
    # 将图像转换为RGB模式
    image = image.convert('RGB')
    
    # 将图像数据转换为一维数组
    pixels = np.array(image).reshape((-1, 3))
    
    # 使用KMeans聚类算法提取主色调
    kmeans = KMeans(n_clusters=num_colors)
    kmeans.fit(pixels)
    palette = kmeans.cluster_centers_
    
    # 将RGB颜色转换为十六进制
    hex_palette = ["#" + "".join([format(int(i), '02x') for i in color]) for color in palette]
    
    # 创建调色板预览图像
    palette_img = Image.new('RGB', (num_colors * 40, 100))
    for i, color in enumerate(palette):
        palette_img.paste(Image.new('RGB', (40, 100), tuple(color.astype(int))), (i * 40, 0))
    
    # 将PIL图像转换为numpy数组以便Gradio显示
    palette_img = np.array(palette_img)
    
    return hex_palette, palette_img

# 创建Gradio界面
iface = gr.Interface(
    fn=extract_color_palette,
    inputs=gr.Image(label="上传图像"),  # 直接使用gr.Image组件
    outputs=[
        gr.Textbox(label="色彩十六进制码", lines=1),  # 显示颜色代码的文本框
        gr.Image(label="调色板预览")  # 显示调色板预览的图像
    ],
    title="KMeans调色板提取器",
    description="拖拽图像到此处获取其调色板。"
)

# 启动Gradio应用
iface.launch()